{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "742ffd26",
   "metadata": {},
   "source": [
    "# PyTorch 和 TVM 之间的 DLPack 集成\n",
    "\n",
    "此测试验证以下内容：\n",
    "1. DLPack 从 PyTorch 到 TVM 的转换\n",
    "2. DLPack 从 TVM 到 PyTorch 的转换\n",
    "3. 转换过程中数据完整性的保存\n",
    "4. DLPack 和 numpy 回退机制之间的功能等价性\n",
    "5. 对不支持的数据类型的错误处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "31a919ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入测试框架和必要的库\n",
    "import pytest\n",
    "import torch\n",
    "import tvm\n",
    "from tvm import relax, tir\n",
    "from tvm.script import relax as R, tir as T\n",
    "from tvm.relax import BasePyModule\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae3888aa",
   "metadata": {},
   "source": [
    "## 测试从 PyTorch 张量到 TVM NDArray 的 DLPack 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "66d93562",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一个 PyTorch 张量作为测试数据\n",
    "pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32)\n",
    "\n",
    "# 使用 DLPack 协议将 PyTorch 张量转换为 TVM NDArray\n",
    "tvm_ndarray = tvm.nd.from_dlpack(pytorch_tensor)\n",
    "\n",
    "# 验证转换结果的类型、形状和数据类型是否正确\n",
    "assert isinstance(tvm_ndarray, tvm.nd.NDArray)\n",
    "assert tvm_ndarray.shape == pytorch_tensor.shape\n",
    "assert str(tvm_ndarray.dtype) == str(pytorch_tensor.dtype).replace(\"torch.\", \"\")\n",
    "\n",
    "# 验证转换后的数据值是否保持一致\n",
    "tvm_numpy = tvm_ndarray.numpy()\n",
    "pytorch_numpy = pytorch_tensor.numpy()\n",
    "np.testing.assert_allclose(tvm_numpy, pytorch_numpy, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adfe57f5",
   "metadata": {},
   "source": [
    "## 测试从 GPU 上的 PyTorch 张量到 TVM NDArray 的 DLPack 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0f3b3cec",
   "metadata": {},
   "outputs": [],
   "source": [
    "if tvm.cuda().exist:\n",
    "    # 在 GPU 上创建一个 PyTorch 张量\n",
    "    pytorch_tensor = torch.tensor(\n",
    "        [1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32, device=\"cuda\"\n",
    "    )\n",
    "\n",
    "    # 使用 DLPack 协议将 GPU 上的 PyTorch 张量转换为 TVM NDArray\n",
    "    tvm_ndarray = tvm.nd.from_dlpack(pytorch_tensor)\n",
    "\n",
    "    # 验证转换结果的类型、形状、数据类型和设备是否正确\n",
    "    assert isinstance(tvm_ndarray, tvm.nd.NDArray)\n",
    "    assert tvm_ndarray.shape == pytorch_tensor.shape\n",
    "    assert str(tvm_ndarray.dtype) == str(pytorch_tensor.dtype).replace(\"torch.\", \"\")\n",
    "    assert str(tvm_ndarray.device) == \"cuda:0\"\n",
    "\n",
    "    # 将数据移动到 CPU 以便进行 numpy 转换和比较\n",
    "    tvm_numpy = tvm_ndarray.numpy()\n",
    "    pytorch_numpy = pytorch_tensor.cpu().numpy()\n",
    "    np.testing.assert_allclose(tvm_numpy, pytorch_numpy, atol=1e-5)\n",
    "else:\n",
    "    # 如果 CUDA 不可用，则跳过此测试\n",
    "    pytest.skip(\"CUDA not available\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f51611d",
   "metadata": {},
   "source": [
    "## 测试从 TVM NDArray 到 PyTorch 张量的 DLPack 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "07c461f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# 创建一个 numpy 数组作为原始数据\n",
    "data = np.array([1.0, 2.0, 3.0, 5.0], dtype=\"float32\")\n",
    "# 将 numpy 数组转换为 TVM NDArray\n",
    "tvm_ndarray = tvm.nd.array(data)\n",
    "\n",
    "# 使用 DLPack 协议将 TVM NDArray 转换为 PyTorch 张量\n",
    "pytorch_tensor = torch.from_dlpack(tvm_ndarray)\n",
    "\n",
    "# 验证转换结果的类型、形状和数据类型是否正确\n",
    "assert isinstance(pytorch_tensor, torch.Tensor)\n",
    "assert pytorch_tensor.shape == tvm_ndarray.shape\n",
    "assert pytorch_tensor.dtype == torch.float32\n",
    "\n",
    "# 验证转换后的数据值是否保持一致\n",
    "tvm_numpy = tvm_ndarray.numpy()\n",
    "pytorch_numpy = pytorch_tensor.numpy()\n",
    "np.testing.assert_allclose(tvm_numpy, pytorch_numpy, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c002483",
   "metadata": {},
   "source": [
    "## 测试从 GPU 上的 TVM NDArray 到 PyTorch 张量的 DLPack 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ebe65ab6",
   "metadata": {},
   "outputs": [],
   "source": [
    "if tvm.cuda().exist:\n",
    "    import numpy as np\n",
    "\n",
    "    # 创建一个 numpy 数组作为原始数据\n",
    "    data = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=\"float32\")\n",
    "    # 将 numpy 数组转换为 GPU 上的 TVM NDArray\n",
    "    tvm_ndarray = tvm.nd.array(data, device=tvm.cuda(0))\n",
    "\n",
    "    # 使用 DLPack 协议将 GPU 上的 TVM NDArray 转换为 PyTorch 张量\n",
    "    pytorch_tensor = torch.from_dlpack(tvm_ndarray)\n",
    "\n",
    "    # 验证转换结果的类型、形状、数据类型和设备是否正确\n",
    "    assert isinstance(pytorch_tensor, torch.Tensor)\n",
    "    assert pytorch_tensor.shape == tvm_ndarray.shape\n",
    "    assert pytorch_tensor.dtype == torch.float32\n",
    "    assert pytorch_tensor.device.type == \"cuda\"\n",
    "\n",
    "    # 将数据移动到 CPU 以便进行 numpy 转换和比较\n",
    "    tvm_numpy = tvm_ndarray.numpy()\n",
    "    pytorch_numpy = pytorch_tensor.cpu().numpy()\n",
    "    np.testing.assert_allclose(tvm_numpy, pytorch_numpy, atol=1e-5)\n",
    "else:\n",
    "    # 如果 CUDA 不可用，则跳过此测试\n",
    "    pytest.skip(\"CUDA not available\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da3930c",
   "metadata": {},
   "source": [
    "## 测试 DLPack 往返转换：PyTorch -> TVM -> PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "70c72a9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建 PyTorch 张量\n",
    "original_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32)\n",
    "\n",
    "# 转换到 TVM\n",
    "tvm_ndarray = tvm.nd.from_dlpack(original_tensor)\n",
    "\n",
    "# 再转换回 PyTorch\n",
    "result_tensor = torch.from_dlpack(tvm_ndarray)\n",
    "\n",
    "# 验证往返转换的数据完整性\n",
    "assert torch.allclose(original_tensor, result_tensor, atol=1e-5)\n",
    "assert original_tensor.dtype == result_tensor.dtype\n",
    "assert original_tensor.shape == result_tensor.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01d0cfc1",
   "metadata": {},
   "source": [
    "## 测试不同数据类型的 DLPack 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e1b445de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义要测试的数据类型对\n",
    "test_types = [\n",
    "    (torch.float32, \"float32\"),\n",
    "    (torch.float64, \"float64\"),\n",
    "    (torch.int32, \"int32\"),\n",
    "    (torch.int64, \"int64\"),\n",
    "]\n",
    "\n",
    "# 对每种数据类型进行转换测试\n",
    "for torch_dtype, tvm_dtype in test_types:\n",
    "    # 创建指定数据类型的 PyTorch 张量\n",
    "    pytorch_tensor = torch.tensor([1, 2, 3], dtype=torch_dtype)\n",
    "\n",
    "    # 转换到 TVM\n",
    "    tvm_ndarray = tvm.nd.from_dlpack(pytorch_tensor)\n",
    "\n",
    "    # 再转换回 PyTorch\n",
    "    result_tensor = torch.from_dlpack(tvm_ndarray)\n",
    "\n",
    "    # 验证转换结果\n",
    "    assert torch.allclose(pytorch_tensor, result_tensor, atol=1e-5)\n",
    "    assert pytorch_tensor.dtype == result_tensor.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e37d4705",
   "metadata": {},
   "source": [
    "## 测试不同形状张量的 DLPack 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2f929793",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义要测试的张量形状\n",
    "test_shapes = [\n",
    "    (1,),           # 标量（1维）\n",
    "    (2, 3),         # 矩阵（2维）\n",
    "    (4, 5, 6),      # 3维张量\n",
    "    (1, 1, 1, 1),   # 4维张量（特殊情况）\n",
    "]\n",
    "\n",
    "# 对每种形状进行转换测试\n",
    "for shape in test_shapes:\n",
    "    # 创建指定形状的随机 PyTorch 张量\n",
    "    pytorch_tensor = torch.randn(shape, dtype=torch.float32)\n",
    "\n",
    "    # 转换到 TVM\n",
    "    tvm_ndarray = tvm.nd.from_dlpack(pytorch_tensor)\n",
    "\n",
    "    # 再转换回 PyTorch\n",
    "    result_tensor = torch.from_dlpack(tvm_ndarray)\n",
    "\n",
    "    # 验证转换结果\n",
    "    assert torch.allclose(pytorch_tensor, result_tensor, atol=1e-5)\n",
    "    assert pytorch_tensor.shape == result_tensor.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c476a11c",
   "metadata": {},
   "source": [
    "## 测试 DLPack 和 numpy 转换是否产生相同的结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "653254f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一个大型 PyTorch 张量\n",
    "size = 1000000\n",
    "pytorch_tensor = torch.randn(size, dtype=torch.float32)\n",
    "\n",
    "# 测试 DLPack 转换\n",
    "tvm_ndarray_dlpack = tvm.nd.from_dlpack(pytorch_tensor)\n",
    "\n",
    "# 测试 numpy 转换\n",
    "numpy_array = pytorch_tensor.detach().cpu().numpy()\n",
    "tvm_ndarray_numpy = tvm.nd.array(numpy_array)\n",
    "\n",
    "# 验证两种方法产生相同的结果\n",
    "result_dlpack = torch.from_dlpack(tvm_ndarray_dlpack)\n",
    "result_numpy = torch.from_numpy(tvm_ndarray_numpy.numpy())\n",
    "assert torch.allclose(result_dlpack, result_numpy, atol=1e-5)\n",
    "\n",
    "# 验证数据完整性\n",
    "assert torch.allclose(result_dlpack, pytorch_tensor, atol=1e-5)\n",
    "assert result_dlpack.shape == pytorch_tensor.shape\n",
    "assert result_dlpack.dtype == pytorch_tensor.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "426c29c0",
   "metadata": {},
   "source": [
    "## 测试 DLPack 对不支持操作的错误处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a2254fd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试非连续张量\n",
    "pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32)\n",
    "non_contiguous = pytorch_tensor[::2]  # 创建一个非连续视图\n",
    "\n",
    "# 这应该可以工作（PyTorch 处理非连续张量）\n",
    "try:\n",
    "    tvm_ndarray = tvm.nd.from_dlpack(non_contiguous)\n",
    "    result_tensor = torch.from_dlpack(tvm_ndarray)\n",
    "    assert torch.allclose(non_contiguous, result_tensor, atol=1e-5)\n",
    "except Exception as e:\n",
    "    # 如果失败，这也是可以接受的\n",
    "    ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a0d6dfc",
   "metadata": {},
   "source": [
    "## 测试在 BasePyModule 上下文中的 DLPack 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c42f3647",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: Failed to compile Relax VM: Module has no function 'vm_load_executable'\n"
     ]
    }
   ],
   "source": [
    "# 创建一个简单的 IRModule\n",
    "@T.prim_func\n",
    "def identity_func(A: T.Buffer((3,), \"float32\"), B: T.Buffer((3,), \"float32\")):\n",
    "    for i in T.grid(3):\n",
    "        B[i] = A[i]\n",
    "\n",
    "# 创建 IRModule 并实例化 BasePyModule\n",
    "ir_mod = tvm.IRModule({\"identity_func\": identity_func})\n",
    "device = tvm.cpu(0)\n",
    "py_mod = BasePyModule(ir_mod, device)\n",
    "\n",
    "# 创建 PyTorch 张量\n",
    "input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)\n",
    "\n",
    "# 调用 TIR 函数（这将触发 DLPack 转换）\n",
    "result = py_mod.call_tir(identity_func, [input_tensor], R.Tensor((3,), \"float32\"))\n",
    "\n",
    "# 验证结果\n",
    "assert isinstance(result, torch.Tensor)\n",
    "assert torch.allclose(result, input_tensor, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc84c72d",
   "metadata": {},
   "source": [
    "## 测试 DLPack 转换保持设备一致性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ef3b607d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试 CPU 张量\n",
    "cpu_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)\n",
    "cpu_tvm = tvm.nd.from_dlpack(cpu_tensor)\n",
    "cpu_result = torch.from_dlpack(cpu_tvm)\n",
    "\n",
    "assert cpu_result.device.type == \"cpu\"\n",
    "assert torch.allclose(cpu_tensor, cpu_result, atol=1e-5)\n",
    "\n",
    "# 注意：GPU 测试需要 CUDA/OpenCL 设置\n",
    "# 这是一个基本测试，确保 CPU 工作正常"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1399189e",
   "metadata": {},
   "source": [
    "## 测试 DLPack 转换在可能的情况下共享内存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "80eb950a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建 PyTorch 张量\n",
    "pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32)\n",
    "\n",
    "# 转换到 TVM\n",
    "tvm_ndarray = tvm.nd.from_dlpack(pytorch_tensor)\n",
    "\n",
    "# 修改原始张量\n",
    "pytorch_tensor[0] = 10.0\n",
    "\n",
    "# 转换回 PyTorch\n",
    "result_tensor = torch.from_dlpack(tvm_ndarray)\n",
    "\n",
    "# 结果应该反映修改（内存共享）\n",
    "assert result_tensor[0] == 10.0\n",
    "assert torch.allclose(pytorch_tensor, result_tensor, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f464f119",
   "metadata": {},
   "source": [
    "## 测试批处理操作中的 DLPack 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f1e4f57a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一批张量\n",
    "batch_size = 10\n",
    "pytorch_tensors = [torch.randn(5, dtype=torch.float32) for _ in range(batch_size)]\n",
    "\n",
    "# 将所有张量转换为 TVM\n",
    "tvm_ndarrays = [tvm.nd.from_dlpack(t) for t in pytorch_tensors]\n",
    "\n",
    "# 将所有张量转换回 PyTorch\n",
    "result_tensors = [torch.from_dlpack(t) for t in tvm_ndarrays]\n",
    "\n",
    "# 验证所有转换\n",
    "for i in range(batch_size):\n",
    "    assert torch.allclose(pytorch_tensors[i], result_tensors[i], atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42a1804b",
   "metadata": {},
   "source": [
    "## 测试 DLPack 在边缘情况下的转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6c3a3303",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 空张量测试\n",
    "empty_tensor = torch.tensor([], dtype=torch.float32)\n",
    "empty_tvm = tvm.nd.from_dlpack(empty_tensor)\n",
    "empty_result = torch.from_dlpack(empty_tvm)\n",
    "\n",
    "assert empty_result.shape == empty_tensor.shape\n",
    "assert empty_result.dtype == empty_tensor.dtype\n",
    "\n",
    "# 单元素张量测试\n",
    "single_tensor = torch.tensor([42.0], dtype=torch.float32)\n",
    "single_tvm = tvm.nd.from_dlpack(single_tensor)\n",
    "single_result = torch.from_dlpack(single_tvm)\n",
    "\n",
    "assert single_result.shape == single_tensor.shape\n",
    "assert single_result[0] == 42.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e441acb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
